# -*- coding: UTF-8 -*-
# *******************************************************************
# File Name: generate_pure_background_sample
# > Author: 04000387
# > Created Time: 2025/1/6 14:11
# *******************************************************************
import textwrap

from PIL import Image, ImageFont, ImageDraw
import uuid
from pathlib import Path
import json
import numpy as np
from utils import get_files, read_files


def _calculate_luminance(_color):
    """计算颜色的亮度"""
    return 0.299 * _color[0] + 0.587 * _color[1] + 0.114 * _color[2]


def _get_contrast_color(background_color):
    """根据背景色选择对比度较高的字体颜色"""
    luminance = _calculate_luminance(background_color)
    if luminance > 128:
        # 背景较亮，选择较暗的字体颜色
        return np.random.randint(0, 127), np.random.randint(0, 127), np.random.randint(0, 127)
    else:
        # 背景较暗，选择较亮的字体颜色
        return np.random.randint(128, 255), np.random.randint(128, 255), np.random.randint(128, 255)


def _generate_font_color(background_color, random_color=True):
    if random_color:
        _color = tuple(np.random.randint(0, 256, 3).tolist())
        return _color
    else:
        return tuple(_get_contrast_color(background_color))


def save_generate(save_path, words, boxes, img_path):
    fp = open(save_path, "a+")
    jsonl = {"img_path": str(img_path), "detail": [{"word": word, "box": boxes[i]} for i, word in enumerate(words)]}
    jsonl = json.dumps(jsonl) + "\n"
    fp.write(jsonl)
    fp.close()


def generate_write_info(text: str,
                        start_pos: tuple[int] = (0, 0),
                        font_size_change_ratio: float = 0.1,
                        img_size: tuple[int] = (224, 224),
                        font_size_min=15
                        ):
    text = text.strip().replace(" ", "").replace("\t", "").replace("\n", "")
    words = [it for it in text]

    boxes = []
    used_words = []
    font_sizes = []

    pos = [start_pos[0], start_pos[1]]
    last_line_height = font_size_min
    while True:
        font_size = font_size_min
        if np.random.rand() < font_size_change_ratio:
            font_size += np.random.randint(5, 80)

        if pos[1] + font_size > img_size[0]:
            break

        if pos[0] + font_size > img_size[1]:
            pos[1] += last_line_height + np.random.randint(-8, 8)
            pos[0] = np.random.randint(2, 5)
            last_line_height = font_size_min
            continue

        used_words.append(np.random.choice(words))
        boxes.append((pos[0], pos[1], font_size, font_size))
        font_sizes.append(font_size)

        pos[0] += font_size + np.random.randint(6, 20)
        pos[1] += np.random.randint(-2, 2)

        last_line_height = max(last_line_height, font_size)

    return used_words, boxes, font_sizes


def generate_pure_background_img(text, name=None, save_path="../generate_dir",
                                 is_show=False,
                                 start_pos=(0, 0),
                                 font_size=25,
                                 line_space=4,
                                 img_size=(224, 224)):
    words, boxes, font_sizes = generate_write_info(text, start_pos, img_size=img_size, font_size_min=font_size)
    background_color = tuple(np.random.randint(0, 256, 3).tolist())

    im = Image.new("RGB", img_size, background_color)
    dr = ImageDraw.Draw(im)
    for i, word in enumerate(words):
        font = ImageFont.truetype("font/simsun.ttc", font_sizes[i])
        box = boxes[i]
        font_color = _generate_font_color(background_color, False)
        dr.text((box[0], box[1]), word, font=font, fill=font_color, spacing=line_space)
    if name is None:
        name = str(uuid.uuid4()) + ".png"

    if len(save_path) != 0:
        img_dir = "img"
        img_path = Path(save_path) / img_dir
        if not img_path.exists():
            img_path.mkdir(parents=True)
        file_save_path = str(img_path / str(name))
        im.save(file_save_path)

        txt_path = Path(save_path) / "generate.jsonl"
        save_generate(txt_path, words, boxes, file_save_path)

    if is_show:
        im.show()

    return words, boxes


def create_img(text, name=None, save_path="../generate_dir",
               is_show=False,
               start_pos=(0, 0),
               font_size=25,
               font_color="#000000",
               background_color=(255, 255, 255),
               line_space=4,
               img_size=(224, 224)
               ):
    img_dir = "img"
    save_path = Path(save_path) / img_dir
    if not save_path.exists():
        save_path.mkdir(parents=True)

    # 计算每张图片最多容纳的字
    line_word_num = img_size[1] // font_size
    line_num = img_size[0] // (font_size + line_space)

    wrapper = textwrap.wrap(text, line_word_num)
    if len(wrapper) >= line_num:
        wrapper = wrapper[:line_num]

    text = "".join(wrapper)

    # 画布颜色
    im = Image.new("RGB", img_size, background_color)
    dr = ImageDraw.Draw(im)
    # 字体样式
    # fontPath = os.path.join("C:\\Windows\\Fonts\\", "simhei.ttf")

    font = ImageFont.truetype("font/simsun.ttc", font_size)
    # 文字颜色
    write_text = "\n".join(wrapper)
    dr.text(start_pos, write_text, font=font, fill=font_color, spacing=line_space)
    if name is None:
        name = str(uuid.uuid4()) + ".png"

    file_save_path = save_path / str(name)
    im.save(file_save_path.absolute())
    if is_show:
        im.show()

    return file_save_path.absolute(), text


def create_txt(text, name, save_path="../generate_dir"):
    txt_dir = "txt"
    save_path = Path(save_path)
    if not save_path.exists():
        save_path.mkdir(parents=True)

    save_path = save_path / "txt.jsonl"

    fp = save_path.open("a+", encoding="utf-8")
    if len(text) == 0:
        text = "这是一张背景图,无文字"

    else:
        text = "这是一张带有文字的图片,图片中的文字如下:" + text
    save_text = json.dumps({"img_file": str(name.absolute()), "content": text}) + "\n"
    fp.write(save_text)
    fp.close()


def generate_init_sample(save_path="../generate_dir/generate", _num=-1):
    _tem_path = Path(save_path).parent

    files = get_files(_tem_path.absolute() / "txt", ".txt")
    for file in files:
        # save_path = Path(save_path) / "sanguo.txt"
        fp = file.open("r", encoding="utf-8")
        lines = fp.readlines()
        fp.close()

        for i, line in enumerate(lines, start=1):
            line = line.strip()

            background_color = tuple(np.random.randint(0, 256, 3).tolist())

            font_color = _generate_font_color(background_color, False)
            img_file, text = create_img(line, is_show=False, save_path=save_path, background_color=background_color,
                                        font_color=font_color, font_size=15)
            create_txt(text, img_file, save_path)

            if 0 < _num <= i:
                break


def batch_generate_pure_background_img(text_path, save_path="../generate_dir/generate", img_size=(224, 224), _num=-1):
    _txt_path = Path(text_path)
    _txt_files = get_files(_txt_path, ".txt")
    texts = read_files(_txt_files)
    num = len(texts)
    if _num > 0:
        num = _num

    for i in range(num):
        generate_pure_background_img(texts[i], save_path=save_path, is_show=False, img_size=img_size)
